from PyCmpltrtok.common import sep
from neo4j import GraphDatabase
from python_nlp.kg.neo4j.config import NEO4J_CONFIG_AGR
from x00100_confirm_location import confirm_location

level_up_or_down_map = {
    'up': 'located in the administrative territorial entity',
    'down': 'contains administrative territorial entity',
}


def level_up_or_down(driver, xbase, up_or_down='up'):
    xupper = []
    with driver.session() as ss:
        cypher = 'match (n:Node{name:%r})-[r:Predicate{name:%r}]->(m:Node) return m.name' % (xbase, level_up_or_down_map[up_or_down], )
        record = ss.run(cypher)
        if record:
            for row in record:
                xupper.append(row[0])

        opposite = list(set(['up', 'down']) - set([up_or_down]))[0]

        cypher = 'match (n:Node{name:%r})<-[r:Predicate{name:%r}]-(m:Node) return m.name' % (xbase, level_up_or_down_map[opposite], )
        record = ss.run(cypher)
        if record:
            for row in record:
                xupper.append(row[0])
    return sorted(set(xupper))  # 去重、排序、返回


def choose(xcandidates):
    if not xcandidates:
        return None

    xlen = len(xcandidates)
    xmax = xlen - 1
    for i, xcan in enumerate(xcandidates):
        print(i, xcan)

    while True:

        print(f'您要选哪一个？（0~{xmax}）：')
        try:
            xnumber = None
            xnumber = int(input().strip())
        except:
            pass
        if xnumber is not None and 0 <= xnumber <= xmax:
            xtarget = xcandidates[xnumber]
            break

    return xtarget


if '__main__' == __name__:

    def _main():
        # 驱动
        driver = GraphDatabase.driver(**NEO4J_CONFIG_AGR)

        while True:
            sep()
            print('请输入查询的地区：')
            xname = input().strip()
            xcandidates = confirm_location(driver, xname)

            if not xcandidates:
                continue

            xlen = len(xcandidates)
            if 1 == xlen:
                xtarget = xcandidates[0]
            else:
                xtarget = choose(xcandidates)

            break

        while True:
            sep(xtarget)
            print("""0 - up
1 - down""")
            while True:
                up_or_down = None
                pool = ['up', 'down']
                try:
                    xinput = int(input().strip())
                    up_or_down = pool[xinput]
                except:
                    pass
                if up_or_down is not None:
                    break
            print(up_or_down, ':')
            xresult = level_up_or_down(driver, xtarget, up_or_down)
            print(xresult)
            if not xresult:
                print('【没有那一级的信息】')
                continue
            if 1 == len(xresult):
                xtarget = xresult[0]
            else:
                xtarget = choose(xresult)

    _main()
